#!/usr/bin/env python3
"""
Robot Motion Editor + Player  (PyBullet + Tkinter)

● 左侧 Treeview 按帧序号展示 13 个关节角度
● 右侧可视化编辑、增删改、上下移动
● CSV 保存 / 加载（motion.csv）
● 播放器：起止帧、帧时长、循环
● 实时将角度写入 PyBullet 仿真
"""

import pybullet as p, pybullet_data
import tkinter as tk
from tkinter import ttk, messagebox
import threading, time, math, csv, os, sys

# ---------- 常量 ----------
JOINT_NAMES   = ['S1','S2','S3','S4','S5','S6','S7','S8','S9','S10','S11']
CSV_PATH      = "U1-3.csv"
DEFAULT_FORCE = 5.0

# ---------- 仿真全局 ----------
robot_id      = None
joint_indices = {}
frames        = []              # list[dict{joint:deg}]

# ---------- PyBullet 仿真线程 ----------
def sim_thread():
    global robot_id, joint_indices
    p.connect(p.GUI)
    p.setAdditionalSearchPath(pybullet_data.getDataPath())
    p.setGravity(0, 0, -9.81)

    p.loadURDF("plane.urdf")

    start_pos = [0, 0, 0.18]

    # start_orn = p.getQuaternionFromEuler([-math.pi / 2, 0, 0])  # Y‑up → Z‑up
    start_orn = p.getQuaternionFromEuler([0, 0, 0])  # Y‑up → Z‑up
    robot_id  = p.loadURDF("RobotPanda3.urdf", start_pos, start_orn)

    # # 设置地面和摩擦
    # plane_id = p.loadURDF("plane.urdf")
    # p.changeDynamics(plane_id, -1, lateralFriction=2.0)
    #
    # start_pos = [0, 0, 0.18]
    # start_orn = p.getQuaternionFromEuler([-math.pi / 2, 0, 0])
    # robot_id = p.loadURDF("PANDA3.urdf", start_pos, start_orn)
    #
    # # 设置机器人各部件的摩擦力
    # num_joints = p.getNumJoints(robot_id)
    # for link_index in range(-1, num_joints):  # 包括 base 和所有 link
    #     p.changeDynamics(robot_id, link_index,
    #                      lateralFriction=0.01,
    #                      spinningFriction=0.1,
    #                      rollingFriction=0.1,
    #                      frictionAnchor=True)

    for i in range(p.getNumJoints(robot_id)):
        name = p.getJointInfo(robot_id, i)[1].decode()
        if name in JOINT_NAMES:
            joint_indices[name] = i

    while True:
        p.stepSimulation()
        time.sleep(1 / 240)

# ---------- 将一帧角度写入仿真 ----------
def apply_frame(frame_dict):
    if robot_id is None:
        return
    for j, deg in frame_dict.items():
        idx = joint_indices.get(j)
        if idx is not None:
            p.setJointMotorControl2(
                robot_id, idx, p.POSITION_CONTROL,
                targetPosition=math.radians(deg),
                force=DEFAULT_FORCE
            )

# ---------- Tk GUI ----------
class MotionGUI(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Robot Motion Editor")
        self.geometry("900x600")
        self.protocol("WM_DELETE_WINDOW", self.on_close)

        # ===== 左侧 Treeview =====
        self.tree = ttk.Treeview(self, columns=JOINT_NAMES, show="headings", height=18)
        for j in JOINT_NAMES:
            self.tree.heading(j, text=j)
            self.tree.column(j, width=55, anchor="center")
        vsb = ttk.Scrollbar(self, orient="vertical", command=self.tree.yview)
        self.tree.configure(yscrollcommand=vsb.set)
        self.tree.bind("<<TreeviewSelect>>", self.on_select)

        # ===== 右侧编辑区 =====
        edit = ttk.Frame(self)
        self.current_idx = None
        self.joint_vars  = {}
        col = 0
        for i, j in enumerate(JOINT_NAMES):
            ttk.Label(edit, text=j).grid(row=i % 7, column=col * 2, sticky="e", padx=2, pady=2)
            var = tk.StringVar(value="0.0")
            self.joint_vars[j] = var
            ttk.Entry(edit, textvariable=var, width=6).grid(row=i % 7, column=col * 2 + 1, padx=2, pady=2)
            if i % 7 == 6:
                col += 1

        # ----- 按钮条 -----
        btns = ttk.Frame(edit); btns.grid(row=8, columnspan=4, pady=5)
        ttk.Button(btns, text="新增帧", command=self.add_frame).grid(row=0, column=0, padx=3)
        ttk.Button(btns, text="插入帧", command=self.insert_frame).grid(row=0, column=1, padx=3)
        ttk.Button(btns, text="删除帧", command=self.delete_frame).grid(row=0, column=2, padx=3)
        ttk.Button(btns, text="上移",   command=lambda: self.move_frame(-1)).grid(row=0, column=3, padx=3)
        ttk.Button(btns, text="下移",   command=lambda: self.move_frame(1)).grid(row=0, column=4, padx=3)
        ttk.Button(btns, text="更新",   command=self.update_frame).grid(row=0, column=5, padx=3)
        ttk.Button(btns, text="保存CSV", command=self.save_csv).grid(row=0, column=6, padx=3)
        ttk.Button(btns, text="加载CSV", command=self.load_csv).grid(row=0, column=7, padx=3)

        # ----- 播放器 -----
        play = ttk.Frame(edit); play.grid(row=9, columnspan=4, pady=5, sticky="w")
        ttk.Label(play, text="起:").pack(side="left")
        self.start_var = tk.StringVar(value="0")
        ttk.Entry(play, width=4, textvariable=self.start_var).pack(side="left")
        ttk.Label(play, text="止:").pack(side="left")
        self.end_var = tk.StringVar(value="")
        ttk.Entry(play, width=4, textvariable=self.end_var).pack(side="left")
        ttk.Label(play, text="帧时(s):").pack(side="left")
        self.dt_var = tk.StringVar(value="1.0")
        ttk.Entry(play, width=4, textvariable=self.dt_var).pack(side="left", padx=3)
        self.loop_var = tk.BooleanVar()
        ttk.Checkbutton(play, text="循环", variable=self.loop_var).pack(side="left")
        ttk.Button(play, text="▶ 播放", command=self.play).pack(side="left", padx=4)
        ttk.Button(play, text="■ 停止", command=self.stop).pack(side="left")

        # ===== 布局 =====
        self.tree.pack(side="left", fill="both", expand=True)
        vsb.pack(side="left", fill="y")
        edit.pack(side="left", fill="y", padx=10, pady=10)

        self.play_flag = False
        self.load_csv()

    # ===== Tree helpers =====
    def refresh_tree(self):
        self.tree.delete(*self.tree.get_children())
        for idx, f in enumerate(frames):
            self.tree.insert("", "end", iid=str(idx), values=[f[j] for j in JOINT_NAMES])
        self.end_var.set(str(max(0, len(frames) - 1)))

    def on_select(self, _):
        sel = self.tree.selection()
        if not sel:
            return
        self.current_idx = int(sel[0])
        f = frames[self.current_idx]
        for j in JOINT_NAMES:
            self.joint_vars[j].set(str(f[j]))

    # ===== Frame ops =====
    def get_entry_frame(self):
        try:
            return {j: float(self.joint_vars[j].get()) for j in JOINT_NAMES}
        except ValueError:
            messagebox.showerror("输入错误", "请确保所有关节角度为数字")
            return None

    def add_frame(self):
        f = self.get_entry_frame()
        if f:
            frames.append(f)
            self.refresh_tree()

    def insert_frame(self):
        if self.current_idx is None:
            return
        f = self.get_entry_frame()
        if f:
            frames.insert(self.current_idx, f)
            self.refresh_tree()

    def delete_frame(self):
        if self.current_idx is not None:
            frames.pop(self.current_idx)
            self.refresh_tree()

    def move_frame(self, step):
        if self.current_idx is None:
            return
        new = self.current_idx + step
        if 0 <= new < len(frames):
            frames[self.current_idx], frames[new] = frames[new], frames[self.current_idx]
            self.refresh_tree()
            self.tree.selection_set(str(new))

    def update_frame(self):
        if self.current_idx is None:
            return
        f = self.get_entry_frame()
        if f:
            frames[self.current_idx] = f
            self.refresh_tree()

    # ===== CSV =====
    def save_csv(self):
        try:
            with open(CSV_PATH, "w", newline="") as f:
                w = csv.writer(f)
                for fr in frames:
                    w.writerow([fr[j] for j in JOINT_NAMES])
            messagebox.showinfo("保存成功", f"已写入 {CSV_PATH}")
        except Exception as e:
            messagebox.showerror("保存失败", str(e))

    def load_csv(self):
        frames.clear()
        if os.path.exists(CSV_PATH):
            with open(CSV_PATH) as f:
                for row in csv.reader(f):
                    if len(row) == len(JOINT_NAMES):
                        frames.append({j: float(v) for j, v in zip(JOINT_NAMES, row)})
        if not frames:
            frames.append({j: 0.0 for j in JOINT_NAMES})
        self.refresh_tree()

    # ===== 播放 =====
    def play(self):
        if self.play_flag:
            return
        try:
            start = int(self.start_var.get())
            end = int(self.end_var.get() or len(frames) - 1)
            dt = float(self.dt_var.get())
        except ValueError:
            messagebox.showerror("输入错误", "起止帧和帧时长需为数字")
            return
        if not (0 <= start <= end < len(frames)):
            messagebox.showerror("范围错误", "起止帧不合法")
            return

        self.play_flag = True

        def _player():
            while self.play_flag:
                for idx in range(start, end + 1):
                    if not self.play_flag:
                        break
                    apply_frame(frames[idx])
                    time.sleep(dt)
                if not self.loop_var.get():
                    break
            self.play_flag = False

        threading.Thread(target=_player, daemon=True).start()

    def stop(self):
        self.play_flag = False

    def on_close(self):
        self.stop()
        self.destroy()
        sys.exit(0)

# ---------- main ----------
if __name__ == "__main__":
    threading.Thread(target=sim_thread, daemon=True).start()
    MotionGUI().mainloop()

